#include "Server.h"

Server::Server(unsigned int start, unsigned int argc, char * const argv[], bool d) : daemon(d)
{
	if (argc-start == 2)
		init(atoi(argv[start+0]), atoi(argv[start+1]));
	else if (argc-start == 1)
		init(atoi(argv[start+0]));
	else
		init();
}

Server::~Server( void )
{
	Dispatch *disconnect = new Dispatch(SERVER_DISCONNECT, "Goodbye!", strlen("Goodbye!"));
	notifyClients( disconnect );
	delete disconnect;
}

void Server::init(void)
{
	init(MAX_CONNECTIONS, TEXT_PORT);
}

void Server::init(unsigned int max)
{
	init(max, TEXT_PORT);
}

void Server::init(unsigned int max, unsigned int port)
{
	max_fd = -1;

	maxConnects = max;
	chat_port = port;

	struct sockaddr_in file_socket_addr;

	/* create a socket
	 * IP protocol family (PF_INET)
	 * TCP protocol (SOCK_STREAM)
	 *****************************************/

	ifprintf("Opening TCP socket...\n");
	if ( (chat_connect_socket = socket(AF_INET, SOCK_STREAM, 0)) == INVALID_SOCKET) return perror("TCP socket");

	/* create a socket
	 * IP protocol family (PF_INET)
	 * TCP protocol (SOCK_STREAM)
	 *****************************************/
    if ( (file_connect_socket = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)) == INVALID_SOCKET) return perror("UDP socket");

	ifprintf("Setting TCP socket options...\n");

#ifdef __WIN32__

	u_long mode = 1;
	ioctlsocket(file_connect_socket, FIONBIO, &mode);

	const char* optval = "1";
	if (setsockopt(chat_connect_socket, SOL_SOCKET, SO_REUSEADDR, optval, sizeof(optval)) == INVALID_SOCKET)
#else

	// Multiplexing isn't workign quite right, so we're making this socket non blocking
	fcntl(file_connect_socket, F_SETFL, O_NONBLOCK);

	int optval=1;
	if (setsockopt(chat_connect_socket, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval)) == INVALID_SOCKET)
#endif
		return perror("setsockopt");

	chat_socket_addr.sin_family = AF_INET;			// host byte order
	chat_socket_addr.sin_port = htons(chat_port);	// short, network byte order
	chat_socket_addr.sin_addr.s_addr = htonl(INADDR_ANY);	// automatically fill with my IP
	memset(&(chat_socket_addr.sin_zero), '\0', 8);	// zero the rest of the struct

	ifprintf("Binding TCP port %d.\n", chat_port);
	if (bind(chat_connect_socket, (struct sockaddr *)&chat_socket_addr, sizeof(struct sockaddr)) == -1) return perror("TCP bind");
	if (listen(chat_connect_socket, maxConnects) == -1) return perror("listen");

	file_socket_addr.sin_family = AF_INET;
	file_socket_addr.sin_port = htons(FILE_PORT);
	file_socket_addr.sin_addr.s_addr = htonl(INADDR_ANY);
	memset(&(file_socket_addr.sin_zero), '\0', 8);	// zero the rest of the struct

	ifprintf("Binding UDP port %d.\n", FILE_PORT);
	if (bind(file_connect_socket, (struct sockaddr*)(&file_socket_addr), sizeof(struct sockaddr)) == -1) return perror("UDP bind");

	ifprintf("Listening for clients...\n");

	// Clear out all socket sets
	FD_ZERO(&client_set);

	// Add TCP server connection socket to set
	FD_SET(chat_connect_socket, &client_set);

	// Add UDP server connection socket to set
	FD_SET(file_connect_socket, &client_set);
}

void Server::dictateClient( ClientConnection *client )
{
	Dispatch d = Dispatch(SERVER_DICTATE_INFO, std::string(1, client->chat_socket) + ":" + client->username + std::string(1, '\t') );
	d.release(client->chat_socket);
}

void Server::acceptConnection( int socket )
{
	ClientConnection *client;

	if (socket == INVALID_SOCKET) return perror("Socket");

	client = new ClientConnection();
	client->chat_socket = socket;
	client->connected = false;
	client->corrupt = 0;
	clients.push_back(client);

	// Add client to master client socket set. (file descriptors)
	FD_SET(socket, &client_set);
	if((int)socket > max_fd) max_fd = socket;

	ifprintf("Met client on socket %u. Awaiting pariticipation request.\n", client->chat_socket);
}

void Server::completeClientConnection( ClientConnection *client, std::string username )
{
	setClientUsername(client, username);
	client->connected = true;

	// TODO: Make sure no users of same name.

	// Let new client know what it's uid and username is.
	dictateClient(client);

	ifprintf("Participation request recieved from %u\n", client->chat_socket);
	ifprintf("I've just met %s. I've got %d client(s) now!\n", client->username.c_str(), clients.size());

	ifprintf("Userlist: ");
	notifyClientsUserList();
}

void Server::notifyClientsUserList( void )
{
	std::string list = listUsers();
	Dispatch *communicationDispatch = new Dispatch(SERVER_USER_LIST, list.c_str(), list.size());
	notifyClients(communicationDispatch);
	delete communicationDispatch;
}

void Server::notifyClients( Dispatch *dispatch, ClientConnection* exclude )
{
	ClientConnection *client;
	// ClientConnection list iterator
	std::list<ClientConnection*>::iterator c_it;

	// Forward all messags (for now) to other clients
	ifprintf("Relaying - %u : %s\n", (unsigned int)dispatch->getFlag(), dispatch->getData().c_str());

	ifprintf("To: ");
	for ( c_it = clients.begin(); c_it != clients.end(); c_it++ )
	{
		client = (*c_it);

		// Do not send ones own message back to self
		if( (exclude != NULL && client == exclude) || !client->connected ) continue;

		// Try to send data to clients socket
		if( dispatch->release( client->chat_socket ) )
			ifprintf("%s,", client->username.c_str());
	}
	ifprintf("\n");
}

void Server::disconnectClient( ClientConnection* client )
{
	Dispatch *notification = new Dispatch(SERVER_DISCONNECT, "Goodbye!");
	notification->release(client->chat_socket);
	// Remove socket from master socket set (Stop listening to it)
	FD_SET(client->chat_socket, &client_set);

#ifdef __WIN32__
	closesocket(client->chat_socket);
#else
	close(client->chat_socket);
#endif

	ifprintf("Disconnecting client %d (%s).\n", client->chat_socket, client->username.c_str());

	// Pass along that same message to other clients
	ifprintf("Notifying users of updated user list");
	notifyClientsUserList();
	delete notification;

}

void Server::setMaxFD(void)
{
	// ClientConnection list iterator
	std::list<ClientConnection*>::iterator c_it;
	max_fd = chat_connect_socket > file_connect_socket ? chat_connect_socket : file_connect_socket;
	for ( c_it = clients.begin(); c_it != clients.end(); c_it++ )
		max_fd = (*c_it)->chat_socket > max_fd ? (*c_it)->chat_socket : max_fd;
}

Server::ClientConnection* Server::findClientByUID( uint16_t c_id )
{
	std::list<ClientConnection*>::iterator c_it;
	for ( c_it = clients.begin(); c_it != clients.end(); c_it++ )
		if ( (*c_it)->chat_socket == c_id ) return (*c_it);
	return NULL;
}

Server::ClientConnection* Server::findClientByUsername( std::string c_name )
{
	std::list<ClientConnection*>::iterator c_it;
	for ( c_it = clients.begin(); c_it != clients.end(); c_it++ )
		if ( (*c_it)->username.compare(c_name) == 0 ) return (*c_it);
	return NULL;
}

std::string Server::listUsers( ClientConnection *exclude )
{
	std::string list;
	ClientConnection *c;

	std::list<ClientConnection*>::iterator c_it;
	for ( c_it = clients.begin(); c_it != clients.end(); c_it++ )
	{
		c = (*c_it);
		// Is this client to be excluded or not connected?
		if (c == exclude || !c->connected) continue;

		// chat_socket is the uid for the client
		list += std::string(1, c->chat_socket) + ":" + c->username + std::string(1, '\t');
	}
	return list;
}

void Server::processFileTransferMessages( ClientConnection *from, Dispatch *d, DispatchFlag flag )
{
	FileTransferHandshake *ft_handshake = new FileTransferHandshake;
	ClientConnection *to;
	char *filename;
	std::string message;

	const char *data = d->getRawData();
	const char *read_ptr = data;

	// Create FileTransferHandshake from data
	memcpy(ft_handshake, read_ptr, sizeof(FileTransferHandshake));
	read_ptr += sizeof(FileTransferHandshake);

	// Find client by UID
	if ( (to = findClientByUID(ft_handshake->client_uid)) != NULL )
	{
		// Reuse FileTransferHandshake as own.
		ft_handshake->client_uid = from->chat_socket;
		ft_handshake->port = from->udp.sin_port;
		ft_handshake->addr = from->udp.sin_addr;

		// Allocate space for the filename
		filename = (char*)malloc(ft_handshake->filename_size+1);
		// Extract filename from buffer.
		memcpy(filename, read_ptr, ft_handshake->filename_size);
		filename[ft_handshake->filename_size] ='\0';
		// advance read_ptr
		read_ptr += ft_handshake->filename_size;


		size_t length = sizeof(FileTransferHandshake)+ft_handshake->filename_size;
		char *buffer = (char*)malloc(length);
		char *write_ptr = buffer;

		memcpy(write_ptr, ft_handshake, sizeof(FileTransferHandshake));
		write_ptr += sizeof(FileTransferHandshake);

		memcpy(write_ptr, filename, ft_handshake->filename_size);


		if (flag == FILE_SEND_ACCEPT)
			message = "has agreed to accept the file";
		else if (flag == FILE_SEND_REJECT)
			message = "has rejected the file transfer";
		else if (flag == FILE_SEND_REQUEST)
			message = "is requesting to send the file";
		else
			message = "error";

		ifprintf("%s(%u) %s '%s' from %s(%u)\n",
			from->username.c_str(),
			from->chat_socket,
			message.c_str(),
			read_ptr,
			to->username.c_str(),
			to->chat_socket);

		/* Forward file transfer acceptance to target client
		 * Format: CUID.FTUID.FN_LENGTH.FILENAME.PORT.IP
		 **************************************************/
		Dispatch d = Dispatch(flag, buffer, length);
		d.release(to->chat_socket);

		// mem cleanup
		free(buffer);
		return;
	}
}

bool Server::setClientUsername( ClientConnection *client, std::string username )
{
	// We need the server to be able to change the clients username.
	// For instance, when we start it up in silent server mode, we're going
	// to supply a default username... but what if someone connects with the same
	// default username? we need to say "no sorry, you have to use default_name_
	// or default_name__ and so on. Like IRC does :D.

	while (findClientByUsername(username)) username += "_";

	// We can assume we did not find another client with the same username.
	client->username = username;

	// Success!
	return true;
}

void Server::handleClientMessage( ClientConnection *client, Dispatch *d )
{
	uint16_t clientid;
	const char *data = d->getRawData();

	data += decode_data(data, &clientid);

	ifprintf("Client Message: %s\n", data);

	if (clientid == client->chat_socket)
		// Notify all clients of outgoing dispatch. (Optional, do not notify 'client')
		notifyClients(d, client);
}

void Server::processCommunicationFromTCP( fd_set readfrom )
{
	ClientConnection *client;

	// Dispatches for incoming and outgoing messages
	Dispatch *incoming;

	struct sockaddr_in si_other;
	socklen_t slen = sizeof(struct sockaddr_in);

	// ClientConnection list iterator
	std::list<ClientConnection*>::iterator c_it;


	if ( FD_ISSET(file_connect_socket, &readfrom) )
	{
		uint16_t clientid;
		char data[sizeof(clientid)];

		if (recvfrom(file_connect_socket, data, sizeof(clientid), 0, (struct sockaddr*)(&si_other), &slen) == -1) perror("UDP recvfrom");
        // The client's public UDP endpoint data is now in si_other.
        // Notice that we're completely ignoring the datagram payload.
        // If we want to support multiple clients inside the same NAT,
        // we'd have clients send their own private UDP endpoints
        // encoded in some way inside the payload, and store those as
        // well.

		decode_data(data, &clientid);

		for ( c_it = clients.begin(); c_it != clients.end(); c_it++ )
		{
			// De-reference the iterator to get the client at position.
			client = *c_it;

			// Check the unique id (chat_socket) of client
			if (client->chat_socket == (unsigned int)clientid )
			{
				printf("Received packet from client %d (%s) at %s:%d\n", clientid, client->username.c_str(), inet_ntoa(si_other.sin_addr), ntohs(si_other.sin_port));

				client->udp = si_other;
				break;
			}
		}
	}

	// For every connected client, process packets
	for ( c_it = clients.begin(); c_it != clients.end(); c_it++ )
	{
		// De-reference the iterator to get the client at position.
		client = *c_it;

		// Is there something to read from the client?
		if( FD_ISSET(client->chat_socket, &readfrom) )
		{
			incoming = new Dispatch();
			incoming->receive(client->chat_socket);

			if (client->connected) switch (incoming->getFlag())
			{
				case CLIENT_DISCONNECT:
					disconnectClient(client);
					clients.erase(c_it++);
					notifyClientsUserList();

					// Reset and find the highest file descriptor
					setMaxFD();

					break;
				case CLIENT_MESSAGE:
					handleClientMessage(client, incoming);
					break;

				case CLIENT_REQUEST_NAME:
					// Attempt to set users name as requested
					if (setClientUsername(client, incoming->getData()))
					{
						// If successful, let client know...
						dictateClient(client);
						// ... and notify all other clients.
						notifyClientsUserList();
					}
					break;

				case FILE_SEND_REQUEST:
				case FILE_SEND_ACCEPT:
				case FILE_SEND_REJECT:
					processFileTransferMessages( client, incoming, incoming->getFlag());
					break;

				case CORRUPT:
				default:
					if (++client->corrupt >= CORRUPT_CLIENT_THRESHOLD)
					{
						ifprintf("(%u) Got too many corrupt packets from %s...\n", client->corrupt, client->username.c_str());
						disconnectClient(client);
						clients.erase(c_it++);

						// Reset and find the highest file descriptor
						setMaxFD();

					} else {
						ifprintf("(%u) Got a packet we didn't understand from %s...\n", client->corrupt, client->username.c_str());
					}
					break;

			// Client has not completed connection protocol
			} else switch (incoming->getFlag())
			{
				case CLIENT_REQUEST_CONNECTION:
					completeClientConnection(client, incoming->getData());
					break;
				default:
					ifprintf("Client has not completed connection protocol. Attempting: %u\n", incoming->getFlag());
					break;
			}

			delete incoming;
		}
	}
}

void Server::run(void)
{

	// File Descriptor set for clients. Gets reset each loop.
	fd_set readfrom;
	FD_ZERO(&readfrom);

	// Socket Address structure for new connections
	struct sockaddr_in new_addr;
	socklen_t sin_size = sizeof(struct sockaddr_in);

	while(1)
	{
		// This acts as the clear and reset for the socket set we're working with
		// since the set gets modified on 'select(...)'
		readfrom = client_set;

		// select() blocks until there is something to read on one of the file descriptors
		if(max_fd != -1 && select(max_fd+1, &readfrom, NULL, NULL, NULL) < 0)
			return perror("select");

		// New connection?
		if(FD_ISSET(chat_connect_socket, &readfrom))
			acceptConnection(accept(chat_connect_socket, (struct sockaddr *)&new_addr, &sin_size));

		processCommunicationFromTCP( readfrom );

	}
}
